import torch
from torch import nn

from .sparse_super_token_encoder import SparseSuperTokenEncoder

from spconv.pytorch import SparseConvTensor

class SparseSuperTokenEncoder_OpenPCDet(nn.Module):
    """a warpper for OpenPCDet"""
    
    def __init__(self, model_cfg, input_channels, voxel_size, grid_size, point_cloud_range):
        super().__init__()
        # global configs
        out_channels = model_cfg.OUT_CHANNELS
        ffn_ratio = model_cfg.FFN_RATIO
        num_heads = model_cfg.NUM_HEADS
        supertoken_iters = model_cfg.SUPERTOKEN_ITERS
        association_cnt = model_cfg.ASSOCIATION_CNT
        temperature = model_cfg.TEMPERATURE
        dropout = model_cfg.DROPOUT
        sta_sampling_mode = model_cfg.STA_CFG.SAMPLING_MODE
        sta_association_mode = model_cfg.STA_CFG.ASSOCIATION_MODE
        # per-layer configs
        encoder_channels = model_cfg.ENCODER_CHANNELS
        encoder_blocks = model_cfg.ENCODER_BLOCKS
        downsample_strides = model_cfg.DOWNSAMPLE_STRIDES
        supertoken_window_sizes = model_cfg.SUPERTOKEN_WINDOW_SIZES
        # fixed configs
        in_channels = input_channels
        sparse_shape = grid_size[::-1] # XXX: inverse order in actual input
        return_middle_feats = False
        return_sparse_output = True
        z_index = 0
        # generated configs
        layer_cfg_base = dict()
        layer_sta_cfgs = []
        layer_sampling_cfgs = []
        layer_association_cfgs = []
        num_layers = len(encoder_blocks)
        for l in range(num_layers):
            layer_sta_cfgs.append(dict(
                num_heads=num_heads,
                num_super_token_iters=supertoken_iters,
                dropout=dropout,
            ))
            sampling_cfgs = []
            association_cfgs = []
            for v in range(2):
                kernel_size = supertoken_window_sizes[l]
                kernel_offset = [v * (k // 2) for k in kernel_size]
                sampling_cfgs.append(dict(
                    mode=sta_sampling_mode,             
                    num_super_tokens_per_batch=1225,
                ))
                association_cfgs.append(dict(
                    mode=sta_association_mode,
                    k=association_cnt,
                ))
            layer_sampling_cfgs.append(sampling_cfgs)
            layer_association_cfgs.append(association_cfgs)
        self.sste = SparseSuperTokenEncoder(
            in_channels, out_channels, encoder_channels, encoder_blocks,
            downsample_strides, ffn_ratio, sparse_shape, 
            layer_cfg_base, layer_sta_cfgs,
            layer_sampling_cfgs, layer_association_cfgs,
            return_middle_feats, return_sparse_output, 
            z_index, temperature
        )
        # other props that OpenPCDet need
        self.num_point_features = in_channels
        self.backbone_channels = out_channels
        
    def forward(self, batch_dict):
        # torch.save(batch_dict['points'], "point_coors.pt")
        feats = batch_dict['voxel_features']
        indices = batch_dict['voxel_coords'].long()
        batch_size = batch_dict['batch_size']
        ret = self.sste(feats, indices, batch_size)
        assert isinstance(ret, SparseConvTensor)
        batch_dict['pillar_features'] = batch_dict['voxel_features'] = ret.features
        batch_dict['voxel_coords'] = ret.indices
        return batch_dict